import pymysql
import pandas as pd
from surprise import SVD, Reader, Dataset
from surprise.model_selection import GridSearchCV
from collections import defaultdict
from surprise.dump import dump, load
from recommand_demo.celery import app



# 从数据库提取数据并保存到CSV文件
def export_data_to_csv():
    conn = pymysql.connect(
        host='localhost',
        port=3306,
        db='movie',
        user='root',
        password='111111',
        charset='utf8'
    )
    cursor = conn.cursor()
    sql = 'SELECT user_id, movie_id, behavior_score FROM behavior'
    cursor.execute(sql)
    data = cursor.fetchall()
    df = pd.DataFrame(data, columns=['user_id', 'movie_id', 'rating'])
    df.to_csv('data.csv', index=False)
    cursor.close()
    conn.close()


# 使用CSV文件进行推荐系统构建和预测
class Recommand_UserCf:
    def __init__(self, csv_path):
        self.csv_path = csv_path
        self.path = './model.plf'

    def trans_to_SVD(self):
        data = pd.read_csv(self.csv_path)
        ratings_reader = Reader(rating_scale=(1, 10))
        train_data = Dataset.load_from_df(data, reader=ratings_reader)
        return train_data

    def build_model(self):
        train_data = self.trans_to_SVD()
        param_grid = {'n_epochs': [10], 'lr_all': [0.02, 0.05], 'reg_all': [0.15, 0.2]}
        gs = GridSearchCV(SVD, param_grid, measures=['rmse', 'mae'], cv=2)
        gs.fit(train_data)
        svd_model = gs.best_estimator['rmse']
        fit_data = train_data.build_full_trainset()
        svd_model.fit(trainset=fit_data)
        return svd_model

    def save_model(self, model, save_path):
        model = self.build_model()
        dump(self.path, algo=model)

    def load_model(self):
        model = load(self.path)
        return model

    def predict_all(self):
        model = self.build_model()
        train_data = self.trans_to_SVD()
        fit_data = train_data.build_full_trainset()
        predicting_set = fit_data.build_anti_testset()
        prediction_list = model.test(predicting_set)
        return prediction_list

    def build_user_prediction(self):
        prediction_list = self.predict_all()
        prediction_dict = defaultdict(list)
        for p in prediction_list:
            prediction_dict[str(p.uid)].append((p.iid, p.est))
        sorted_prediction_dict = {}
        for user_id, item_ratings in prediction_dict.items():
            item_ratings.sort(key=lambda x: x[1], reverse=True)
            sorted_prediction_dict[user_id] = item_ratings
        return sorted_prediction_dict

    def tok_k_recommand(self, user_id, k, times):
        predicted_ratings = self.build_user_prediction()
        position = times * k + 1
        if predicted_ratings:
            movie_list = predicted_ratings[user_id][position:position + k]
        return movie_list
@app.task
def Buildmain(csv_path):
    recommender = Recommand_UserCf(csv_path)
    model = recommender.build_model()
    # 在这里可以将模型保存到文件中，如果需要的话
    return model

if __name__ == '__main__':
    # 导出数据到CSV文件
    export_data_to_csv()
    # 使用CSV文件进行推荐
    recommender = Recommand_UserCf('data.csv')
    top_k = recommender.tok_k_recommand('1', 6, 1)
    print(top_k)
